﻿//     Copyright (c) Microsoft Corporation.  All rights reserved.

using System;
using System.Collections.Generic;
using System.Configuration;
using System.Security.Claims;
using System.Security.Cryptography;
using System.Web;
using System.Web.Mvc;
using System.Web.Routing;
using Gateway.Console.Services.Azure;
using Microsoft.Gateway.Utils;
using Microsoft.IdentityModel.Clients.ActiveDirectory;

namespace Gateway.Console.Services
{
    public enum ResourceSelectors
    {
        GraphApi,
        RDFE,
    }

    public static class OAuth2AuthenticationExtensions
    {
        class SelectorParameters
        {
            public string SessionKey { get; set; }
            public string Resource { get; set; }
        }

        class AuthenticationInfo
        {
            public string AuthorizationCode { get; set; }
            public bool AuthorizationCodeFailed { get; set; }
            public AuthenticationResult AccessToken { get; set; }
            public string ErrorCode { get; set; }
            public string ErrorDescription { get; set; }
            public Uri RedirectUri { get; set; }
            public string InitiatingController { get; set; }
            public string InitiatingAction { get; set; }
            public bool IsProvisioningPermission { get; set; }
            public string AuthCodeState { get; set; }
        }

        private static readonly Dictionary<ResourceSelectors, SelectorParameters> Selectors = new Dictionary<ResourceSelectors, SelectorParameters>
        {
            {ResourceSelectors.GraphApi,    new SelectorParameters { SessionKey = "_GraphAuth", Resource = "00000002-0000-0000-c000-000000000000"}},
            {ResourceSelectors.RDFE,        new SelectorParameters { SessionKey = "_RdfeAuth", Resource = AzureEnvironment.RdfeManagementEndpoint}},
        };

        const string TenantIdClaimType          = "http://schemas.microsoft.com/identity/claims/tenantid";
        const string ReplyController            = "Home";
        const string ReplyAction                = "AuthorizationCode";
        static readonly string AppPrincipalId   = Configuration.ClientID;
        static readonly string AppKey           = Configuration.ClientKey;

        public static string LoginUrlBase()
        {
            string tenantId = "common";
            if (ClaimsPrincipal.Current != null)
            {
                var tenantClaim = ClaimsPrincipal.Current.FindFirst(TenantIdClaimType);
                if (tenantClaim != null)
                {
                    tenantId = ClaimsPrincipal.Current.FindFirst(TenantIdClaimType).Value;
                }
            }
            return String.Format("{0}{1}", AzureEnvironment.AadLoginEndpoint, HttpUtility.UrlEncode(tenantId));
        }

        private static AuthenticationInfo GetAuthInfo(this HttpSessionStateBase src, ResourceSelectors selector)
        {
            return src.GetAuthInfo(selector, null);
        }

        private static AuthenticationInfo GetAuthInfo(this HttpSessionStateBase src, ResourceSelectors selector, string requestState, bool createIfNotExist = true)
        {
            var info = src[Selectors[selector].SessionKey] as AuthenticationInfo;
            if (info == null && createIfNotExist)
            {
                info = new AuthenticationInfo();
                src[Selectors[selector].SessionKey] = info;
            }
            if (!String.IsNullOrWhiteSpace(requestState) && info != null)
            {
                if (src[requestState] == null)
                {
                    src[requestState] = info;
                    info.AuthCodeState = requestState;
                }
            }
            return info;
        }

        private static AuthenticationInfo GetAuthInfo(this HttpSessionStateBase src, string requestState)
        {
            return src[requestState] as AuthenticationInfo;
        }

        public static bool IsValidAuthorizationState(this HttpSessionStateBase src, string requestState)
        {
            return src.GetAuthInfo(requestState) != null;
        }

        public static void ResetOAuthInformation(this HttpSessionStateBase src, ResourceSelectors selector)
        {
            var authInfo = src.GetAuthInfo(selector, null, false);
            if (authInfo != null && !String.IsNullOrWhiteSpace(authInfo.AuthCodeState))
            {
                src.Remove(authInfo.AuthCodeState);
        }
            src.Remove(Selectors[selector].SessionKey);
        }

        public static void ResetAllOAuthInformation(this HttpSessionStateBase src)
        {
            foreach (ResourceSelectors selector in Enum.GetValues(typeof(ResourceSelectors)))
            {
                src.ResetOAuthInformation(selector);
            }
        }

        public static bool HasAuthorizationCode(this HttpSessionStateBase src, ResourceSelectors selector)
        {
            return !String.IsNullOrWhiteSpace(src.GetAuthInfo(selector).AuthorizationCode);
        }

        public static bool HasAuthorizationCodeFailed(this HttpSessionStateBase src, ResourceSelectors selector)
        {
            return src.GetAuthInfo(selector).AuthorizationCodeFailed;
        }

        public static bool HasAccessToken(this HttpSessionStateBase src, ResourceSelectors selector)
        {
            return src.GetAuthInfo(selector).AccessToken != null;
        }

        public static bool IsProvisioningPermission(this HttpSessionStateBase src, ResourceSelectors selector)
        {
            return src.GetAuthInfo(selector).IsProvisioningPermission;
        }

        public static void PermissionIsBeingProvisioned(this HttpSessionStateBase src, ResourceSelectors selector)
        {
            src.GetAuthInfo(selector).IsProvisioningPermission = true;
        }

        public static bool GetAuthorizationCodeFailureDetails(this HttpSessionStateBase src, ResourceSelectors selector, out string errorCode, out string errorDescription)
        {
            errorCode = String.Empty;
            errorDescription = String.Empty;
            var authInfo = src.GetAuthInfo(selector);
            if (authInfo.AuthorizationCodeFailed)
            {
                errorCode = authInfo.ErrorCode;
                errorDescription = authInfo.ErrorDescription;
                return true;
            }
            return false;
        }

        public static void AuthorizationCodeFailed(this HttpSessionStateBase src, string requestState, string errorCode, string errorDescription)
        {
            var authInfo = src.GetAuthInfo(requestState);
            if (authInfo != null)
            {
                authInfo.AuthorizationCodeFailed = true;
                authInfo.ErrorCode = errorCode;
                authInfo.ErrorDescription = errorDescription;
                authInfo.AuthorizationCode = String.Empty;
                authInfo.AccessToken = null;
            }
        }

        public static ActionResult GetAuthorizationCode(this HttpSessionStateBase src, ResourceSelectors selector, string initiatingController, string initiatingAction = null)
        {
            string replyProtocol = HttpContext.Current.Request.Url.IsDefaultPort ? Uri.UriSchemeHttps : HttpContext.Current.Request.Url.Scheme;
            string redirectUri = UrlHelper.GenerateUrl("Default", ReplyAction, ReplyController, replyProtocol, HttpContext.Current.Request.Url.Host, "",
                null, RouteTable.Routes, HttpContext.Current.Request.RequestContext, false);
            return src.GetAuthorizationCode(selector, new Uri(redirectUri), initiatingController, initiatingAction);
        }

        public static ActionResult GetAuthorizationCode(this HttpSessionStateBase src, ResourceSelectors selector, Uri redirectUri, string initiatingController, string initiatingAction = null)
        {
            var authCodeState = GenerateCsrfToken();
            string authorizationUrl = string.Format("{0}/oauth2/authorize?api-version=1.0&response_type=code&client_id={1}&resource={2}&redirect_uri={3}&state={4}",
                LoginUrlBase(),
                HttpUtility.UrlEncode(AppPrincipalId),
                HttpUtility.UrlEncode(Selectors[selector].Resource),
                HttpUtility.UrlEncode(redirectUri.ToString()),
                authCodeState);
            var authInfo = src.GetAuthInfo(selector, authCodeState);
            authInfo.AuthorizationCode = String.Empty;
            authInfo.AccessToken = null;
            authInfo.RedirectUri = redirectUri;
            authInfo.InitiatingController = initiatingController;
            authInfo.InitiatingAction = initiatingAction;
            return new RedirectResult(authorizationUrl);
        }

        private static string GenerateCsrfToken()
        {
            var rnd = new Random();
            var buffer = new byte[20];
            rnd.NextBytes(buffer);
            string retval = HttpServerUtility.UrlTokenEncode(buffer);
            Array.Clear(buffer, 0, buffer.Length);
            return retval;
        }

        public static AuthenticationResult GetAccessToken(this HttpSessionStateBase src, string requestState, string authorizationCode = null)
        {
            var authInfo = src.GetAuthInfo(requestState);
            if (authInfo != null)
            {
                return src.GetAccessToken(authInfo, authorizationCode);
            }
            return null;
        }

        public static AuthenticationResult GetAccessToken(this HttpSessionStateBase src, ResourceSelectors selector, string authorizationCode = null)
        {
            return src.GetAccessToken(src.GetAuthInfo(selector), authorizationCode);
        }

        private static AuthenticationResult GetAccessToken(this HttpSessionStateBase src, AuthenticationInfo authInfo, string authorizationCode)
        {
            if (String.IsNullOrWhiteSpace(authorizationCode))
            {
                if (authInfo.AccessToken != null)
                {
                    if (authInfo.AccessToken.ExpiresOn > DateTime.UtcNow)
                    {
                        return authInfo.AccessToken;
                    }
                    else if (!String.IsNullOrWhiteSpace(authInfo.AccessToken.RefreshToken))
                    {
                        authInfo.AccessToken = GetAccessTokenFromRefreshToken(authInfo.AccessToken.RefreshToken);
                        return authInfo.AccessToken;
                    }
                }
                authorizationCode = authInfo.AuthorizationCode;
                authInfo.AccessToken = null;
            }
            else
            {
                authInfo.AuthorizationCode = authorizationCode;
                authInfo.AccessToken = null;
                authInfo.AuthorizationCodeFailed = false;
                authInfo.ErrorCode = String.Empty;
                authInfo.ErrorDescription = String.Empty;
                authInfo.IsProvisioningPermission = false;
            }
            var redirectUri = authInfo.RedirectUri;
            if (!String.IsNullOrWhiteSpace(authorizationCode) && redirectUri != null)
            {
                var ctx = new AuthenticationContext(LoginUrlBase());
                var appCreds = new ClientCredential(AppPrincipalId, AppKey);
                var result = ctx.AcquireTokenByAuthorizationCode(authorizationCode, redirectUri, appCreds);
                authInfo.AccessToken = result;
                return result;
            }
            return null;
        }

        public static AuthenticationResult GetAccessTokenFromRefreshToken(string refreshToken)
        {
            var ctx = new AuthenticationContext(LoginUrlBase());
            var appCreds = new ClientCredential(AppPrincipalId, AppKey);
            return ctx.AcquireTokenByRefreshToken(refreshToken, AppPrincipalId, appCreds);
        }

        public static ActionResult RedirectToAuthenticationInitiator(this HttpSessionStateBase src, string requestState)
        {
            var authInfo = src.GetAuthInfo(requestState);
            if (authInfo != null)
            {
                // Reset the auth state now to prevent replay attack
                authInfo.AuthCodeState = null;
                src.Remove(requestState);
                return new RedirectToRouteResult(new RouteValueDictionary() 
                {
                    { "controller", authInfo.InitiatingController},
                    { "action", authInfo.InitiatingAction },
                });
            }
            return new HttpUnauthorizedResult();
        }

        public static AadGraphApi GetGraphApi(this HttpSessionStateBase src)
        {
            var graphCreds = src.GetAccessToken(ResourceSelectors.GraphApi);
            if (graphCreds != null)
            {
                return new AadGraphApi(graphCreds.TenantId, graphCreds);
            }
            throw new InvalidOperationException("GraphAPI credentials are not valid");
        }
    }
}
